import jieba
import pandas as pd
from config import *

# 加载配置
config = Config()


# 定义函数读取指定文件数据，并转为fasttext输入格式
def process_data(datapath, processed_datapath, is_char=True):
    # 读取原始数据（假设文件是制表符分隔，无表头）
    df = pd.read_csv(
        datapath,
        sep='\t',
        header=None,
        names=['text', 'label'],
        encoding='utf-8',
        on_bad_lines='skip'  # 自动跳过格式错误的行（pandas >= 1.3）
    )
    # 清理文本中的换行符、回车符（防止写入时断行）
    df['text'] = df['text'].astype(str).str.replace(r'[\n\r\t]+', ' ', regex=True).str.strip()
    # 转换标签列为整数
    df['label'] = pd.to_numeric(df['label'], errors='coerce')
    # 删除标签无法转为整数的行
    df = df.dropna(subset=['label'])
    df['label'] = df['label'].astype(int)

    # 转换为 fasttext 格式
    def format_line(row):
        text = row['text']
        label_idx = row['label']
        with open(config.class_doc_path, 'r', encoding='utf-8') as f:
            id2class = {i: line.strip() for i, line in enumerate(f)}
        if label_idx not in id2class:
            return None  # 或 raise ValueError

        if is_char:
            text_split = " ".join(list(str(text)))
        else:
            text_split = " ".join(jieba.lcut(str(text)))

        return f"__label__{label_idx} {text_split}\n"

    # 应用格式化
    with open(processed_datapath, 'w', encoding='utf-8') as fw:
        for _, row in df.iterrows():
            line = format_line(row)
            if line is not None:
                fw.write(line)


if __name__ == '__main__':
    # 字符级别处理
    process_data(config.train_datapath, config.process_train_datapath_char, is_char=True)
    process_data(config.test_datapath, config.process_test_datapath_char, is_char=True)
    process_data(config.dev_datapath, config.process_dev_datapath_char, is_char=True)

    # 分词级别处理
    process_data(config.train_datapath, config.process_train_datapath_word, is_char=False)
    process_data(config.test_datapath, config.process_test_datapath_word, is_char=False)
    process_data(config.dev_datapath, config.process_dev_datapath_word, is_char=False)
